
import matplotlib.pyplot as plt

from dataset.mnist import load_mnist
from common.optimizer import *
from common.util import smooth_curve
from common.simple_conv_net import SimpleConvNet
from common.trainer import Trainer

# 0.读入MNIST数据
(x_train, t_train), (x_test, t_test) = load_mnist(flatten=False)

# 减少学习数据
x_train = x_train[:12000]
t_train = t_train[:12000]

# 设置是否使用Dropout及比例
max_epochs = 20

network = SimpleConvNet(input_dim=(1, 28, 28),
                        conv_param={'filter_num':30, 'filter_size':5, 'pad':0, 'stride':1},
                        hidden_size=100, output_size=10, weight_init_std=0.01)
trainer = Trainer(network, x_train, t_train, x_test, t_test,
                  epochs=max_epochs, mini_batch_size=100,
                  optimizer='Adam', optimizer_param={'lr': 0.001},
                  evaluate_sample_num_per_epoch=1000)

# 开始训练
trainer.train()
# 保存训练结果
network.save_params("params.pkl")
print("Saved Network Parameters!")

# 绘制图像
markers = {"train": "o", "test": "s"}
x = np.arange(max_epochs)
plt.plot(x, smooth_curve(trainer.train_acc_list), marker='o', label='train', markevery=2)
plt.plot(x, smooth_curve(trainer.test_acc_list), marker='s', label='test', markevery=2)
plt.xlabel("epochs")
plt.ylabel("accuracy")
plt.ylim(0, 1.0)
plt.legend(loc='lower right')
plt.show()